package com.at.combine9;

import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.api.java.tuple.Tuple3;
import org.apache.flink.streaming.api.datastream.ConnectedStreams;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.co.CoProcessFunction;
import org.apache.flink.util.Collector;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * TODO 连接两条流，输出能根据id匹配上的数据（类似inner join效果）
 *
 * @author huangchao E-mail:fengquan8866@163.com
 * @version 创建时间：2024/9/25 18:32
 */
public class ConnectKeybyDemo3 {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(2);

        DataStreamSource<Tuple2<Integer, String>> source1 = env.fromElements(
                Tuple2.of(1, "a1"),
                Tuple2.of(1, "a2"),
                Tuple2.of(2, "b"),
                Tuple2.of(3, "c")
        );

        DataStreamSource<Tuple3<Integer, String, Integer>> source2 = env.fromElements(
                Tuple3.of(1, "aa1", 1),
                Tuple3.of(1, "aa2", 2),
                Tuple3.of(2, "bb", 1),
                Tuple3.of(3, "cc", 1)
        );

        ConnectedStreams<Tuple2<Integer, String>, Tuple3<Integer, String, Integer>> connect = source1.connect(source2);

        // 多并行度下，需要根据 关联条件进行 keyby，才能保证 key相同的数据到一起去，才能匹配上
        ConnectedStreams<Tuple2<Integer, String>, Tuple3<Integer, String, Integer>> connectKeyby = connect.keyBy(s1 -> s1.f0, s2 -> s2.f0);

        /**
         * 实现互相匹配的效果：  两条流，不一定谁的数据先来
         * 1、每条流，有数据来，存到一个变量中
         *    hashmap
         *    => key=id, 第一个值
         *    =》 value=List<数据>
         * 2、每条流有数据来的时候，除了存变量中，去另一条流存的变量中，查找是否匹配上
         */
        SingleOutputStreamOperator<String> process = connectKeyby.process(
                new CoProcessFunction<Tuple2<Integer, String>, Tuple3<Integer, String, Integer>, String>() {
                    // 定义hashmap，用来存数据
                    Map<Integer, List<Tuple2<Integer, String>>> s1Cache = new HashMap<>();
                    Map<Integer, List<Tuple3<Integer, String, Integer>>> s2Cache = new HashMap<>();

                    @Override
                    public void processElement1(Tuple2<Integer, String> value, CoProcessFunction<Tuple2<Integer, String>, Tuple3<Integer, String, Integer>, String>.Context ctx, Collector<String> out) throws Exception {
                        Integer id = value.f0;
                        // TODO 1. s1的数据来了，就存到变量中
                        List<Tuple2<Integer, String>> tuple2s = s1Cache.computeIfAbsent(id, k -> new ArrayList<>());
                        tuple2s.add(value);
                        // TODO 2. 去 s2Cache 中查找是否有id能匹配上
                        if (s2Cache.containsKey(id)) {
                            for (Tuple3<Integer, String, Integer> s2Element : s2Cache.get(id)) {
                                out.collect("s1:" + value + "<===1====>" + "s2:" + s2Element);
                            }
                        }
                    }

                    @Override
                    public void processElement2(Tuple3<Integer, String, Integer> value, CoProcessFunction<Tuple2<Integer, String>, Tuple3<Integer, String, Integer>, String>.Context ctx, Collector<String> out) throws Exception {
                        Integer id = value.f0;
                        // TODO 1. s1的数据来了，就存到变量中
                        List<Tuple3<Integer, String, Integer>> tuple2s = s2Cache.computeIfAbsent(id, k -> new ArrayList<>());
                        tuple2s.add(value);
                        // TODO 2. 去 s2Cache 中查找是否有id能匹配上
                        if (s1Cache.containsKey(id)) {
                            for (Tuple2<Integer, String> s1Element : s1Cache.get(id)) {
                                out.collect("s1:" + s1Element + "<===2====>" + "s2:" + value);
                            }
                        }
                    }
                }
        );
        process.print();

        env.execute();
    }
}
